import os
import cv2
import json
from tqdm import tqdm


def yolo_segment_to_labelme(yolo_dataset_path, labelme_dataset_path):
    """
    将 YOLO 分割格式的数据集转换为 LabelMe 格式的数据集

    :param yolo_dataset_path: YOLO 分割格式数据集的路径，包含 images 和 labels 文件夹
    :param labelme_dataset_path: 转换后 LabelMe 格式数据集的保存路径
    """
    yolo_images_path = os.path.join(yolo_dataset_path, "images")
    yolo_labels_path = os.path.join(yolo_dataset_path, "labels")
    os.makedirs(labelme_dataset_path, exist_ok=True)

    image_files = [f for f in os.listdir(yolo_images_path) if f.endswith(('.jpg', '.png'))]
    for image_file in tqdm(image_files):
        image_path = os.path.join(yolo_images_path, image_file)
        label_file = image_file.replace('.jpg', '.txt').replace('.png', '.txt')
        label_path = os.path.join(yolo_labels_path, label_file)

        img = cv2.imread(image_path)
        height, width, _ = img.shape

        labelme_data = {
            "version": "5.0.1",
            "flags": {},
            "shapes": [],
            "imagePath": image_file,
            "imageData": None,
            "imageHeight": height,
            "imageWidth": width
        }

        if os.path.exists(label_path):
            with open(label_path, 'r') as f:
                lines = f.readlines()
                for line in lines:
                    parts = line.strip().split()
                    if len(parts) == 0:
                        continue
                    class_index = int(parts[0])
                    points_str = parts[1:]
                    points = []
                    for i in range(0, len(points_str), 2):
                        x = float(points_str[i]) * width
                        y = float(points_str[i + 1]) * height
                        points.append([x, y])

                    shape = {
                        "label": str(class_index),
                        "points": points,
                        "group_id": None,
                        "shape_type": "polygon",
                        "flags": {}
                    }
                    labelme_data["shapes"].append(shape)

        labelme_json_path = os.path.join(labelme_dataset_path,
                                         image_file.replace('.jpg', '.json').replace('.png', '.json'))
        with open(labelme_json_path, 'w') as f:
            json.dump(labelme_data, f, indent=4)


# 示例调用
yolo_dataset_path = r"C:\Users\kang_\Desktop\workspace\pig_pose\yolo\train"  # 替换为你的 YOLO 分割数据集路径
labelme_dataset_path = r"C:\Users\kang_\Desktop\workspace\pig_pose\yolo\train\images"  # 替换为你希望保存 LabelMe 格式数据集的路径
yolo_segment_to_labelme(yolo_dataset_path, labelme_dataset_path)
